package com.gitee.wsl.jvm.okhttp.callback


import com.gitee.wsl.ext.okio.base.CallState
import com.gitee.wsl.ext.okio.source.CountingSource
import com.gitee.wsl.jvm.okhttp.base.StatusCode
import com.gitee.wsl.jvm.okhttp.request.CountingRequestBody
import com.gitee.wsl.jvm.okhttp.response.mayGetSinglePartContentRange
import com.gitee.wsl.jvm.okhttp.response.mayGetTotalLengthOfUnsatisfiedRange
import kotlinx.coroutines.channels.SendChannel
import okhttp3.Call
import okhttp3.Callback
import okhttp3.Response
import okhttp3.ResponseBody.Companion.asResponseBody
import okio.IOException
import okio.buffer


/**
 * OkHttp [Callback] for [callWithCounting].
 */
class CountingCallback<T>(
    private val countingRequestBody: CountingRequestBody?,
    private val countDownload: Boolean,
    private val handlePartialResponse: Boolean,
    private val responseHandler: (Response) -> T,
    private val channel: SendChannel<CallState<T>>
) : Callback {

    override fun onResponse(call: Call, response: Response) {
        runAndCloseChannel {
            val finalState = try {
                response.use { processResponse(it) }
            } catch (e: IOException) {
                CallState.Failure(e)
            }
            // Since the channel is CONFLATED, `trySend()` will always succeed unless the channel is
            // cancelled.
            channel.trySend(finalState)
        }
    }

    override fun onFailure(call: Call, e: IOException) {
        runAndCloseChannel {
            // Since the channel is CONFLATED, `trySend()` will always succeed unless the channel is
            // cancelled.
            channel.trySend(CallState.Failure(e))
        }
    }

    private fun processResponse(response: Response): CallState.Success<T> {
        val responseBuilderForHandler = response.newBuilder()

        val request = response.request
        val requestBody = request.body
        if (requestBody is CountingRequestBody) {
            // Restore the original request body.
            val restoredRequest = request.newBuilder()
                .method(request.method, requestBody.delegate)
                .build()
            responseBuilderForHandler.request(restoredRequest)
        }

        val countingSource = mayCreateCountingSource(response)
        if (countingSource != null) {
            val originalBody = checkNotNull(response.body)
            val countingResponseBody = countingSource.buffer().asResponseBody(originalBody.contentType(), originalBody.contentLength())
            responseBuilderForHandler.body(countingResponseBody)
        }

        val data = responseHandler(responseBuilderForHandler.build())

        val bytesUploaded = countingRequestBody?.bytesUploaded ?: -1
        val bytesDownloaded = when {
            countingSource != null -> countingSource.bytesDownloaded
            countDownload && handlePartialResponse &&
                    response.code == StatusCode.RANGE_NOT_SATISFIABLE ->
                response.mayGetTotalLengthOfUnsatisfiedRange() ?: -1
            else -> -1
        }
        return CallState.Success(data, bytesUploaded, bytesDownloaded)
    }

    private fun mayCreateCountingSource(response: Response): CountingSource? {
        if (!countDownload || !response.isSuccessful) return null

        val originalBody = checkNotNull(response.body)
        val originalSource = originalBody.source()
        val contentRange =
            if (handlePartialResponse && response.code == StatusCode.PARTIAL_CONTENT) {
                response.mayGetSinglePartContentRange()
            } else null
        return if (contentRange != null) {
            CountingSource(originalSource, contentRange.start, contentRange.totalLength, channel)
        } else {
            CountingSource(originalSource, 0L, originalBody.contentLength(), channel)
        }
    }

    private inline fun runAndCloseChannel(block: () -> Unit) {
        var closeCause: Throwable? = null
        try {
            block()
        } catch (e: Throwable) {
            closeCause = e
        } finally {
            channel.close(closeCause)
        }
    }
}